from ultralytics import YOLO
import os
import json
import yaml
from pathlib import Path
import torch
from datetime import datetime
import cv2
import numpy as np

# 可爱标识
星星 = "☆"
叉叉 = "×"
圈圈 = "○"
箭头 = "→"
问号 = "？"

# 参数模板预设
PARAM_TEMPLATES = {
    "1": {  # 快速模式
        "name": "快速",
        "epochs": 10,
        "imgsz": 320,
        "batch": -1,
        "patience": 20,
        "workers": 8
    },
    "2": {  # 高精度模式
        "name": "高精度",
        "epochs": 300,
        "imgsz": 640,
        "batch": 0.7,
        "patience": 100,
        "workers": 8
    },
    "3": {  # 独角兽模式
        "name": "unicorn",
        "epochs": 1000,
        "imgsz": 640,
        "batch": 0.7,
        "patience": 300,
        "workers": 8
    },
    "4": {  # 自定义模式
        "name": "自定义",
        "epochs": 100,
        "imgsz": 640,
        "batch": 16,
        "patience": 100,
        "workers": 8
    }
}

# 数据增强选项说明
AUGMENT_OPTIONS = {
    "1": {"name": "色调调整（hsv_h）", "param": "hsv_h", "min": 0.0, "max": 0.1, "default": 0.015, "help": "调整图片颜色的色调，值越大变化越明显～"},
    "2": {"name": "饱和度调整（hsv_s）", "param": "hsv_s", "min": 0.0, "max": 1.0, "default": 0.7, "help": "调整颜色的鲜艳程度，0是灰色，1是原图饱和度～"},
    "3": {"name": "亮度调整（hsv_v）", "param": "hsv_v", "min": 0.0, "max": 1.0, "default": 0.4, "help": "调整图片明暗，0是黑色，1是原图亮度～"},
    "4": {"name": "水平翻转（fliplr）", "param": "fliplr", "min": 0.0, "max": 1.0, "default": 0.5, "help": "图片左右翻转的概率，0.5就是50%概率翻转～"},
    "5": {"name": "垂直翻转（flipud）", "param": "flipud", "min": 0.0, "max": 1.0, "default": 0.0, "help": "图片上下翻转的概率，适合无上下方向的场景～"},
    "6": {"name": "旋转角度（degrees）", "param": "degrees", "min": 0.0, "max": 90.0, "default": 0.0, "help": "随机旋转图片的角度范围，比如30就是-30°到30°～"},
    "7": {"name": "平移（translate）", "param": "translate", "min": 0.0, "max": 0.5, "default": 0.1, "help": "随机平移图片的比例，0.1就是上下左右移动10%～"},
    "8": {"name": "缩放（scale）", "param": "scale", "min": 0.0, "max": 2.0, "default": 0.5, "help": "随机缩放的比例，0.5就是缩小到50%或放大到150%～"},
    "9": {"name": "剪切（shear）", "param": "shear", "min": 0.0, "max": 30.0, "default": 0.0, "help": "随机剪切图片的角度，会让图片产生斜切效果～"},
    "10": {"name": "马赛克增强（mosaic）", "param": "mosaic", "min": 0.0, "max": 1.0, "default": 1.0, "help": "将4张图拼成1张的概率，增加场景复杂度～"},
    "11": {"name": "随机透视变换（perspective）", "param": "perspective", "min": 0.0, "max": 0.001, "default": 0.0, "help": "对图像进行随机透视变换，增强模型理解三维空间中物体的能力。"},
    "12": {"name": "四图合一（mosaic）", "param": "mosaic", "min": 0.0, "max": 1.0, "default": 1.0, "help": "将四幅训练图像合成一幅，模拟不同的场景构成和物体互动。对复杂场景的理解非常有效。"},
    "13": {"name": "二图合一（mixup）", "param": "mixup", "min": 0.0, "max": 1.0, "default": 0.0, "help": "混合两幅图像及其标签，创建合成图像。通过引入标签噪声和视觉变化，增强模型的泛化能力。"},
    "14": {"name": "加对象实例（copy_paste）", "param": "copy_paste", "min": 0.0, "max": 1.0, "default": 0.0, "help": "仅限分割。在图像中复制和粘贴对象，以增加对象实例。"},
    "15": {"name": "随机擦除特征（erasing）", "param": "erasing", "min": 0.0, "max": 0.9, "default": 0.4, "help": "仅分类。在训练过程中随机擦除图像区域，以鼓励模型专注于不太明显的特征。"},
}

def 加载历史配置():
    配置文件 = Path.home() / ".yolo_trainer_config.json"
    if 配置文件.exists():
        with open(配置文件, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

def 保存配置(配置):
    配置文件 = Path.home() / ".yolo_trainer_config.json"
    try:
        with open(配置文件, "w", encoding="utf-8") as f:
            json.dump(配置, f, indent=2, ensure_ascii=False)
        print(f"{圈圈} 配置存好啦～ (◍•ᴗ•◍)")
    except Exception as e:
        print(f"{叉叉} 保存配置失败了呢：{str(e)}")

def 选文件(提示语, 后缀=None, 帮助信息="", 历史路径=""):
    while True:
        完整提示 = 提示语
        if 历史路径:
            完整提示 += f"[上次选的是：{历史路径}，按回车直接用哦～] "
        
        路径 = input(完整提示).strip()
        if not 路径 and 历史路径:
            路径 = 历史路径
            print(f"{星星} 就用上次的路径啦～：{路径}")
            return Path(路径).resolve()
        
        if 路径 == "?":
            print(f"\n{问号} 帮助：{帮助信息}\n")
            continue
        if not 路径:
            print(f"{叉叉} 路径不能为空哦～ 再输一次嘛～")
            continue
        
        路径 = Path(路径).resolve()
        if not 路径.exists():
            print(f"{叉叉} 找不到这个文件呢：{路径}")
            continue
        if not 路径.is_file():
            print(f"{叉叉} 这不是文件哦～ 要选文件才行呢～")
            continue
        if 后缀 and not 路径.suffix == 后缀:
            print(f"{叉叉} 要选{后缀}格式的文件哦～ 再看看嘛～")
            continue
        
        print(f"{星星} 找到文件啦～ 就用这个：{路径}")
        return 路径

def 选保存父文件夹(帮助信息="", 历史父路径=""):
    """选择保存结果的父文件夹（子文件夹会按时间自动创建）"""
    while True:
        完整提示 = "请输入训练结果的父文件夹路径（子文件夹会按时间自动创建哦～）："
        if 历史父路径:
            完整提示 += f"[上次选的是：{历史父路径}，按回车直接用哦～] "
        
        父路径 = input(完整提示).strip()
        # 回车直接用历史父路径
        if not 父路径 and 历史父路径:
            父路径 = 历史父路径
            父路径 = Path(父路径).resolve()
            print(f"{星星} 就用上次的父文件夹啦～：{父路径}")
            return 父路径
        
        if 父路径 == "?":
            print(f"\n{问号} 帮助：{帮助信息}\n")
            continue
        if not 父路径:
            print(f"{叉叉} 父文件夹路径不能为空哦～ 再输一次嘛～")
            continue
        
        父路径 = Path(父路径).resolve()
        父父目录 = 父路径.parent  # 父文件夹的上级目录
        if not 父父目录.exists():
            print(f"{叉叉} 上级目录不存在呢：{父父目录}")
            continue
        if 父路径.is_file():
            print(f"{叉叉} 这是文件哦～ 要选文件夹才行呢～")
            continue
        
        # 确保父文件夹存在
        if not 父路径.exists():
            try:
                父路径.mkdir(parents=True, exist_ok=True)
                print(f"{圈圈} 新建了父文件夹哦：{父路径}")
            except Exception as e:
                print(f"{叉叉} 建父文件夹失败了：{str(e)}")
                continue
        else:
            print(f"{星星} 找到父文件夹啦～：{父路径}")
        
        return 父路径

def 创建时间子文件夹(父文件夹):
    """在父文件夹下创建按时间命名的子文件夹（格式：年月日_时分秒）"""
    时间字符串 = datetime.now().strftime("%Y%m%d_%H%M%S")  # 例如：20240807_153045
    子文件夹 = 父文件夹 / 时间字符串
    子文件夹.mkdir(exist_ok=True)  # 确保子文件夹存在
    print(f"{圈圈} 自动创建子文件夹啦～：{子文件夹}（按训练开始时间命名的哦～）")
    return 子文件夹

def 输整数(提示语, 最小=None, 默认值=None, 帮助信息=""):
    while True:
        输入 = input(提示语).strip()
        if 输入 == "?":
            print(f"\n{问号} 帮助：{帮助信息}\n")
            continue
        if not 输入 and 默认值 is not None:
            return 默认值
        try:
            数值 = int(输入)
            if 最小 is not None and 数值 < 最小:
                print(f"{叉叉} 要大于等于{最小}哦～ 再输一次嘛～")
                continue
            return 数值
        except ValueError:
            print(f"{叉叉} 要输数字呀～ 比如1、10这样的～")

def 输小数(提示语, 最小=None, 最大=None, 默认值=None, 帮助信息=""):
    while True:
        输入 = input(提示语).strip()
        if 输入 == "?":
            print(f"\n{问号} 帮助：{帮助信息}\n")
            continue
        if not 输入 and 默认值 is not None:
            return 默认值
        try:
            数值 = float(输入)
            if 最小 is not None and 数值 < 最小:
                print(f"{叉叉} 要大于等于{最小}哦～ 再输一次嘛～")
                continue
            if 最大 is not None and 数值 > 最大:
                print(f"{叉叉} 要小于等于{最大}哦～ 再输一次嘛～")
                continue
            return 数值
        except ValueError:
            print(f"{叉叉} 要输数字呀～ 比如0.5、1.2这样的～")

def 检查GPU可用性():
    if torch.cuda.is_available() and torch.cuda.device_count() > 0:
        return True, f"发现{torch.cuda.device_count()}个GPU哦～ 可以用GPU训练呢～"
    else:
        return False, "没找到可用的GPU呢～ 只能用CPU训练啦～ 可能会有点慢哦～"

def 选择数据增强():
    print("\n===== 数据增强设置～ 好丰富的选项呢～ =====")
    print("○ 这里可以选要增强的方式哦～ 输入序号就能配置对应的参数啦～")
    print("○ 想跳过某个选项就不选它，没有选的会按照默认参数增强～")
    print("○ 所有配置完后输入'0'结束～")
    
    for idx, info in AUGMENT_OPTIONS.items():
        print(f"{idx}. {info['name']}（默认值：{info['default']}，范围：{info['min']}-{info['max']}）")
    
    aug_params = {}
    while True:
        选择 = input("\n请输入要配置的增强序号（0结束）：").strip()
        if 选择 == "0":
            print(f"{星星} 数据增强配置好啦～")
            break
        if 选择 not in AUGMENT_OPTIONS:
            print(f"{叉叉} 要输1-{len(AUGMENT_OPTIONS)}之间的序号哦～ 再试试嘛～")
            continue
        
        aug_info = AUGMENT_OPTIONS[选择]
        print(f"\n{圈圈} 现在配置「{aug_info['name']}」～")
        print(f"💡 说明：{aug_info['help']}")
        
        数值 = 输小数(
            f"请输入{aug_info['name']}的值（{aug_info['min']}-{aug_info['max']}，默认{aug_info['default']}）：",
            最小=aug_info['min'],
            最大=aug_info['max'],
            默认值=aug_info['default'],
            帮助信息=aug_info['help']
        )
        aug_params[aug_info['param']] = 数值
        print(f"{星星} 已设置{aug_info['name']}为：{数值}")
    
    return aug_params

def 分析数据集(数据路径):
    """分析数据集并返回统计信息"""
    print(f"\n{箭头} 正在分析数据集... 稍等一下下～")
    
    try:
        with open(数据路径, "r", encoding="utf-8") as f:
            数据 = yaml.safe_load(f)
    except Exception as e:
        print(f"{叉叉} 解析数据集文件失败：{str(e)}")
        return None
    
    # 获取数据集目录
    数据目录 = Path(数据路径).parent
    
    # 统计训练集和验证集大小
    训练图像数量 = 0
    验证图像数量 = 0
    
    # 检查训练集
    训练路径 = Path(数据["train"])
    if not 训练路径.is_absolute():
        训练路径 = (数据目录 / 训练路径).resolve()
    
    if 训练路径.exists():
        # 假设图像文件在images目录下
        训练图像目录 = 训练路径 / "images"
        if 训练图像目录.exists():
            训练图像数量 = len([f for f in 训练图像目录.iterdir() if f.suffix in ['.jpg', '.jpeg', '.png']])
        else:
            # 直接统计训练路径下的图像文件
            训练图像数量 = len([f for f in 训练路径.iterdir() if f.suffix in ['.jpg', '.jpeg', '.png']])
    
    # 检查验证集
    验证路径 = Path(数据["val"])
    if not 验证路径.is_absolute():
        验证路径 = (数据目录 / 验证路径).resolve()
    
    if 验证路径.exists():
        # 假设图像文件在images目录下
        验证图像目录 = 验证路径 / "images"
        if 验证图像目录.exists():
            验证图像数量 = len([f for f in 验证图像目录.iterdir() if f.suffix in ['.jpg', '.jpeg', '.png']])
        else:
            # 直接统计验证路径下的图像文件
            验证图像数量 = len([f for f in 验证路径.iterdir() if f.suffix in ['.jpg', '.jpeg', '.png']])
    
    # 类别数量
    类别数量 = len(数据.get("names", []))
    
    # 采样几个图像计算平均分辨率
    平均宽度 = 0
    平均高度 = 0
    采样数量 = min(10, 训练图像数量)
    
    if 训练图像数量 > 0 and 训练图像目录.exists():
        图像文件 = list(训练图像目录.iterdir())[:采样数量]
        总宽度 = 0
        总高度 = 0
        
        for img_file in 图像文件:
            try:
                img = cv2.imread(str(img_file))
                if img is not None:
                    高度, 宽度 = img.shape[:2]
                    总宽度 += 宽度
                    总高度 += 高度
            except Exception as e:
                print(f"{叉叉} 读取图像{img_file.name}失败：{str(e)}")
                continue
        
        if len(图像文件) > 0:
            平均宽度 = 总宽度 / len(图像文件)
            平均高度 = 总高度 / len(图像文件)
    
    # 构建统计信息
    统计信息 = {
        "训练图像数量": 训练图像数量,
        "验证图像数量": 验证图像数量,
        "总图像数量": 训练图像数量 + 验证图像数量,
        "类别数量": 类别数量,
        "平均宽度": round(平均宽度, 2),
        "平均高度": round(平均高度, 2)
    }
    
    # 打印统计信息
    总图像数量 = 训练图像数量 + 验证图像数量
    print(f"""
☆ 数据集分析结果～
    训练图像数量: {训练图像数量}
    验证图像数量: {验证图像数量}
    总图像数量: {总图像数量}
    类别数量: {类别数量}
""")
    
    if 平均宽度 > 0 and 平均高度 > 0:
        print(f"    平均分辨率: {平均宽度}x{平均高度}")
    else:
        print("    无法计算平均分辨率")
    
    return 统计信息

# 添加训练建议函数
def generate_training_suggestion(train_count, val_count, total_count, class_count, model_name):
    suggestion = []
    
    # 根据数据集大小和类别数量推荐迁移学习策略
    if total_count < 1000:
        strategy = 1  # 固定特征提取器
        strategy_desc = "固定特征提取器（只训练检测头层）"
        freeze_layers = 15
    elif total_count < 5000:
        strategy = 3  # 两阶段训练
        strategy_desc = "两阶段训练（先固定特征提取器，再微调所有层）"
        freeze_layers = 10
        unfreeze_epochs = max(20, total_count // 100)
    else:
        strategy = 2  # 微调所有层
        strategy_desc = "微调所有层（特征提取层使用较小学习率）"
        freeze_layers = 0
    
    # 单类和多类的调整
    if class_count == 1:
        suggestion.append("📌 单类别检测任务，建议适当增加训练轮次以提高精度")
    elif class_count > 10:
        suggestion.append("📌 多类别检测任务，建议增加数据增强强度")
    
    # 根据模型大小调整建议
    if 'n' in model_name.lower() or 's' in model_name.lower():
        suggestion.append("📌 轻量级模型，建议使用较小的批量大小(4-8)和学习率(0.0001-0.001)")
    elif 'm' in model_name.lower():
        suggestion.append("📌 中等规模模型，建议批量大小(8-16)和学习率(0.001-0.005)")
    else:
        suggestion.append("📌 大型模型，建议批量大小(16-32)和学习率(0.005-0.01)")
    
    # 添加策略建议
    suggestion.append(f"🔍 推荐迁移学习策略: {strategy} ({strategy_desc})")
    
    if strategy == 1:
        suggestion.append(f"🧊 建议冻结层数: {freeze_layers}")
    elif strategy == 3:
        suggestion.append(f"🧊 第一阶段建议冻结层数: {freeze_layers}")
        suggestion.append(f"🔓 建议解冻轮次: {unfreeze_epochs}")
        suggestion.append("📉 微调阶段特征提取层学习率比例建议: 0.1")
    
    # 建议训练轮次
    base_epochs = 50
    if total_count < 1000:
        epochs = base_epochs
    elif total_count < 5000:
        epochs = base_epochs * 2
    else:
        epochs = base_epochs * 3
    suggestion.append(f"⏱️ 建议训练总轮次: {epochs}")
    
    # 建议图片尺寸
    if total_count < 1000:
        imgsz = 640
    else:
        imgsz = 1024 if 'x' in model_name.lower() else 640
    suggestion.append(f"🖼️ 建议图片尺寸: {imgsz}")
    
    # 建议批量大小
    if total_count < 1000:
        batch = 4
    elif total_count < 5000:
        batch = 8
    else:
        batch = 16
    suggestion.append(f"📦 建议批量大小: {batch}")
    
    return '\n'.join(suggestion), strategy, freeze_layers, unfreeze_epochs if strategy == 3 else None

# 加载历史配置
def 保存配置(配置):
    配置文件 = Path.home() / ".yolo_trainer_config.json"
    try:
        with open(配置文件, "w", encoding="utf-8") as f:
            json.dump(配置, f, indent=2, ensure_ascii=False)
        print(f"{圈圈} 配置存好啦～ (◍•ᴗ•◍)")
    except Exception as e:
        print(f"{叉叉} 保存配置失败了呢：{str(e)}")

def 选文件(提示语, 后缀=None, 帮助信息="", 历史路径=""):
    while True:
        完整提示 = 提示语
        if 历史路径:
            完整提示 += f"[上次选的是：{历史路径}，按回车直接用哦～] "
        
        路径 = input(完整提示).strip()
        if not 路径 and 历史路径:
            路径 = 历史路径
            print(f"{星星} 就用上次的路径啦～：{路径}")
            return Path(路径).resolve()
        
        if 路径 == "?":
            print(f"\n{问号} 帮助：{帮助信息}\n")
            continue
        if not 路径:
            print(f"{叉叉} 路径不能为空哦～ 再输一次嘛～")
            continue
        
        路径 = Path(路径).resolve()
        if not 路径.exists():
            print(f"{叉叉} 找不到这个文件呢：{路径}")
            continue
        if not 路径.is_file():
            print(f"{叉叉} 这不是文件哦～ 要选文件才行呢～")
            continue
        if 后缀 and not 路径.suffix == 后缀:
            print(f"{叉叉} 要选{后缀}格式的文件哦～ 再看看嘛～")
            continue
        
        print(f"{星星} 找到文件啦～ 就用这个：{路径}")
        return 路径

def 选保存父文件夹(帮助信息="", 历史父路径=""):
    """选择保存结果的父文件夹（子文件夹会按时间自动创建）"""
    while True:
        完整提示 = "请输入训练结果的父文件夹路径（子文件夹会按时间自动创建哦～）："
        if 历史父路径:
            完整提示 += f"[上次选的是：{历史父路径}，按回车直接用哦～] "
        
        父路径 = input(完整提示).strip()
        # 回车直接用历史父路径
        if not 父路径 and 历史父路径:
            父路径 = 历史父路径
            父路径 = Path(父路径).resolve()
            print(f"{星星} 就用上次的父文件夹啦～：{父路径}")
            return 父路径
        
        if 父路径 == "?":
            print(f"\n{问号} 帮助：{帮助信息}\n")
            continue
        if not 父路径:
            print(f"{叉叉} 父文件夹路径不能为空哦～ 再输一次嘛～")
            continue
        
        父路径 = Path(父路径).resolve()
        父父目录 = 父路径.parent  # 父文件夹的上级目录
        if not 父父目录.exists():
            print(f"{叉叉} 上级目录不存在呢：{父父目录}")
            continue
        if 父路径.is_file():
            print(f"{叉叉} 这是文件哦～ 要选文件夹才行呢～")
            continue
        
        # 确保父文件夹存在
        if not 父路径.exists():
            try:
                父路径.mkdir(parents=True, exist_ok=True)
                print(f"{圈圈} 新建了父文件夹哦：{父路径}")
            except Exception as e:
                print(f"{叉叉} 建父文件夹失败了：{str(e)}")
                continue
        else:
            print(f"{星星} 找到父文件夹啦～：{父路径}")
        
        return 父路径

def 创建时间子文件夹(父文件夹):
    """在父文件夹下创建按时间命名的子文件夹（格式：年月日_时分秒）"""
    时间字符串 = datetime.now().strftime("%Y%m%d_%H%M%S")  # 例如：20240807_153045
    子文件夹 = 父文件夹 / 时间字符串
    子文件夹.mkdir(exist_ok=True)  # 确保子文件夹存在
    print(f"{圈圈} 自动创建子文件夹啦～：{子文件夹}（按训练开始时间命名的哦～）")
    return 子文件夹

def 输整数(提示语, 最小=None, 默认值=None, 帮助信息=""):
    while True:
        输入 = input(提示语).strip()
        if 输入 == "?":
            print(f"\n{问号} 帮助：{帮助信息}\n")
            continue
        if not 输入 and 默认值 is not None:
            return 默认值
        try:
            数值 = int(输入)
            if 最小 is not None and 数值 < 最小:
                print(f"{叉叉} 要大于等于{最小}哦～ 再输一次嘛～")
                continue
            return 数值
        except ValueError:
            print(f"{叉叉} 要输数字呀～ 比如1、10这样的～")

def 输小数(提示语, 最小=None, 最大=None, 默认值=None, 帮助信息=""):
    while True:
        输入 = input(提示语).strip()
        if 输入 == "?":
            print(f"\n{问号} 帮助：{帮助信息}\n")
            continue
        if not 输入 and 默认值 is not None:
            return 默认值
        try:
            数值 = float(输入)
            if 最小 is not None and 数值 < 最小:
                print(f"{叉叉} 要大于等于{最小}哦～ 再输一次嘛～")
                continue
            if 最大 is not None and 数值 > 最大:
                print(f"{叉叉} 要小于等于{最大}哦～ 再输一次嘛～")
                continue
            return 数值
        except ValueError:
            print(f"{叉叉} 要输数字呀～ 比如0.5、1.2这样的～")

def 检查GPU可用性():
    if torch.cuda.is_available() and torch.cuda.device_count() > 0:
        return True, f"发现{torch.cuda.device_count()}个GPU哦～ 可以用GPU训练呢～"
    else:
        return False, "没找到可用的GPU呢～ 只能用CPU训练啦～ 可能会有点慢哦～"

def 选择数据增强():
    print("\n===== 数据增强设置～ 好丰富的选项呢～ =====")
    print("○ 这里可以选要增强的方式哦～ 输入序号就能配置对应的参数啦～")
    print("○ 想跳过某个选项就不选它，没有选的会按照默认参数增强～")
    print("○ 所有配置完后输入'0'结束～")
    
    for idx, info in AUGMENT_OPTIONS.items():
        print(f"{idx}. {info['name']}（默认值：{info['default']}，范围：{info['min']}-{info['max']}）")
    
    aug_params = {}
    while True:
        选择 = input("\n请输入要配置的增强序号（0结束）：").strip()
        if 选择 == "0":
            print(f"{星星} 数据增强配置好啦～")
            break
        if 选择 not in AUGMENT_OPTIONS:
            print(f"{叉叉} 要输1-{len(AUGMENT_OPTIONS)}之间的序号哦～ 再试试嘛～")
            continue
        
        aug_info = AUGMENT_OPTIONS[选择]
        print(f"\n{圈圈} 现在配置「{aug_info['name']}」～")
        print(f"💡 说明：{aug_info['help']}")
        
        数值 = 输小数(
            f"请输入{aug_info['name']}的值（{aug_info['min']}-{aug_info['max']}，默认{aug_info['default']}）：",
            最小=aug_info['min'],
            最大=aug_info['max'],
            默认值=aug_info['default'],
            帮助信息=aug_info['help']
        )
        aug_params[aug_info['param']] = 数值
        print(f"{星星} 已设置{aug_info['name']}为：{数值}")
    
    return aug_params

def main():
    print("☆～ 欢迎使用Unicorn的YOLO训练小助手～ ☆")
    
    # 提前检测GPU并提示
    gpu可用, gpu信息 = 检查GPU可用性()
    print(f"\n{圈圈} 硬件检查：{gpu信息}")
    
    历史配置 = 加载历史配置()
    历史父文件夹 = 历史配置.get("save_parent_dir", "")  # 读取上次的父文件夹路径

    # 断点续训选择
    续训选择 = input("\n要继续上次没训完的模型吗？(y/n，输?看说明)：").strip().lower()
    if 续训选择 == "?":
        print("\n？ 断点训练就是从上次保存的last.pt继续哦,配置上上一次的选择～ 要至少训完1个epoch才行呢～\n")
        续训选择 = input("要继续上次的训练吗？(y/n)：").strip().lower()
    续训 = 续训选择 == 'y'

    # 模型路径选择
    模型帮助 = "模型是.pt格式的哦～ 续训就选last.pt，新训练就选预训练模型（比如yolo11n.pt）～"
    历史模型路径 = ""
    if 续训:
        模型提示 = "请输入last.pt的路径（输?看帮助）："
        历史模型路径 = 历史配置.get("last_model", "")
    else:
        模型提示 = "请输入预训练模型的路径（.pt，输?看帮助）："
        历史模型路径 = 历史配置.get("pretrained_model", "")
    
    模型路径 = 选文件(模型提示, ".pt", 模型帮助, 历史模型路径)

    # 数据集选择
    数据帮助 = "数据集是.yaml格式的哦～ 里面要写train和val的路径，还有类别名字（names）呢～"
    历史数据路径 = 历史配置.get("dataset", "")
    数据路径 = None
    while 数据路径 is None:
        数据提示 = "请输入数据集配置文件的路径（.yaml，输?看帮助）："
        候选路径 = 选文件(数据提示, ".yaml", 数据帮助, 历史数据路径)
        
        print(f"\n{箭头} 正在检查数据集哦～ 稍等一下下～")
        try:
            with open(候选路径, "r", encoding="utf-8") as f:
                数据 = yaml.safe_load(f)
            有效 = True
            信息 = "数据集检查通过～ 没问题啦～"
            for 字段 in ["train", "val", "names"]:
                if 字段 not in 数据:
                    有效 = False
                    信息 = f"数据集文件里少了'{字段}'呢～ 要加上才行哦～"
                    break
            if 有效:
                数据目录 = Path(候选路径).parent
                for 分机 in ["train", "val"]:
                    分机路径 = Path(数据[分机])
                    if not 分机路径.is_absolute():
                        分机路径 = (数据目录 / 分机路径).resolve()
                    if not 分机路径.exists():
                        有效 = False
                        信息 = f"{分机}的路径不对哦：{分机路径}\n（提示：可以写相对当前yaml文件的路径呢～）"
                        break
        except Exception as e:
            有效 = False
            信息 = f"解析失败了：{str(e)}"
        
        if not 有效:
            print(f"{叉叉} 数据集有点问题呢：{信息}")
            print(f"{圈圈} 再选一个正确的数据集文件好不好呀～\n")
            历史数据路径 = ""
        else:
            print(f"{星星} {信息}")
            数据路径 = 候选路径
            
            # 分析数据集并提供迁移学习建议
    数据集统计 = 分析数据集(数据路径)
    if 数据集统计:
        train_count = 数据集统计['训练图像数量']
        val_count = 数据集统计['验证图像数量']
        total_count = 数据集统计['总图像数量']
        class_count = 数据集统计['类别数量']
        model_name = 模型路径.name

        # 生成训练建议
        training_suggestion, default_strategy, default_freeze, default_unfreeze = generate_training_suggestion(
            train_count, val_count, total_count, class_count, model_name)
        print(f"\n☆ 训练建议～\n{training_suggestion}")
    else:
        default_strategy = 3
        default_freeze = 10
        default_unfreeze = 50

    # 加载模型
    try:
        模型 = YOLO(模型路径)
        model_name = 模型路径.name
        print(f"{星星} 模型加载好啦～ 是{model_name}呢～")

        # 打印模型结构信息
        print(f"{圈圈} 模型结构：{模型.model.__class__.__name__}")
        print(f"{圈圈} 总层数：{len(list(模型.model.parameters()))}层")
    except Exception as e:
        print(f"{叉叉} 模型加载失败了：{str(e)}")
        return
    
    # 添加迁移学习策略选择
    print("\n===== 迁移学习设置～ =====")
    print("○ 迁移学习策略：")
    print("1. 固定特征提取器：只训练检测头层")
    print("2. 微调：训练所有层，但特征提取层使用较小学习率")
    print("3. 两阶段训练：先固定特征提取器，再微调所有层")

    迁移学习策略 = input(f"请选择迁移学习策略（1/2/3，默认{default_strategy}）：").strip() or str(default_strategy)
    while 迁移学习策略 not in ["1", "2", "3"]:
        迁移学习策略 = input("请输入1、2或3：").strip() or str(default_strategy)
    
    冻结层数 = 0
    微调学习率比例 = 0.1  # 微调阶段特征提取层学习率比例
    解冻轮次 = 0
    
    if 迁移学习策略 == "1":
        print("\n{圈圈} 策略1: 固定特征提取器，只训练检测头层～")
        冻结层数 = 输整数(
            "要冻结多少层特征提取层？(默认10，输?看帮助)：",
            最小=0,
            默认值=10,
            帮助信息="YOLO模型通常有多个特征提取层，冻结层数越多，训练速度越快～"
        )
    elif 迁移学习策略 == "2":
        print("\n{圈圈} 策略2: 微调所有层，特征提取层使用较小学习率～")
        冻结层数 = 0  # 不冻结任何层
        微调学习率比例 = 输小数(
            "特征提取层学习率比例（相对于检测头层，默认0.1）：",
            最小=0.01,
            最大=1.0,
            默认值=0.1,
            帮助信息="设置特征提取层学习率为检测头层的多少倍，小于1表示特征提取层学习率更低～"
        )
    else:
        print("\n{圈圈} 策略3: 两阶段训练～")
        冻结层数 = 输整数(
            f"第一阶段要冻结多少层特征提取层？(默认{default_freeze})：",
            最小=0,
            默认值=default_freeze,
        )
        解冻轮次 = 输整数(
            f"训练多少轮后解冻所有层进行微调？(默认{default_unfreeze})：",
            最小=1,
            默认值=default_unfreeze,
        )
        微调学习率比例 = 输小数(
            "微调阶段特征提取层学习率比例（默认0.1）：",
            最小=0.01,
            最大=1.0,
            默认值=0.1,
        )
    
    # 参数模板选择
    print("\n===== 要选个参数模板吗～ =====")
    print("○ 输入序号就能选哦～ 选了之后会直接用模板的参数呢～")
    print("1. 快速：训练10轮，图片320尺寸，自动适配GPU～ 适合快速测试效果～")
    print("2. 高精度：训练300轮，图片640尺寸，用70% GPU～ 适合正式训练呢～")
    print("3. unicorn（独角兽模式）：训练1000轮，超有耐心～ 适合追求极致精度的小伙伴～")
    print("4. 自定义：可以自己调所有参数，想怎么设就怎么设～")
    
    模板选择 = ""
    while 模板选择 not in PARAM_TEMPLATES.keys():
        模板选择 = input("请输入序号（1/2/3/4，默认4）：").strip() or "4"
        if 模板选择 not in PARAM_TEMPLATES.keys():
            print(f"{叉叉} 要输1、2、3、4哦～ 再试试嘛～")
    
    模板 = PARAM_TEMPLATES[模板选择]
    print(f"{星星} 选了「{模板['name']}」模式～ 真乖～")

    # 训练参数配置（只有自定义模式需要手动设置）
    if 模板选择 == "4":
        print("\n===== 来设置训练参数吧～ =====")
        轮次 = 输整数(
            f"要训练多少轮呢（默认{模板['epochs']}，输?看帮助）：",
            最小=1,
            默认值=模板['epochs'],
            帮助信息="就是把所有数据训多少遍哦～ 少了可能学不会，多了会很慢呢～"
        )
        图片尺寸 = 输整数(
            f"图片要缩放到多大呢（默认{模板['imgsz']}，输?看帮助）：",
            最小=32,
            默认值=模板['imgsz'],
            帮助信息="正方形的尺寸哦～ 太小看不清，太大训起来慢～ 320-1280之间比较好～"
        )
        批量 = 输小数(
            f"每次喂多少张图片呢（默认{模板['batch']}，输?看帮助）：",
            最小=-1,
            默认值=模板['batch'],
            帮助信息="可以输数字（比如16）、-1（自动适配）或0.7（70% GPU）～ 太大了装不下哦～"
        )
        if isinstance(批量, float) and 批量.is_integer():
            批量 = int(批量)
        早停耐心 = 输整数(
            f"多久没进步就停下呢（默认{模板['patience']}，输?看帮助）：",
            最小=0,
            默认值=模板['patience'],
            帮助信息="比如设20，就是20轮没进步就提前停～ 省时间呢～ 0的话就一直训完～"
        )
        线程数 = 输整数(
            f"数据加载用多少线程呢（默认{模板['workers']}，输?看帮助）：",
            最小=0,
            默认值=模板['workers'],
            帮助信息="和电脑CPU核心数差不多就好～ 比如8核就设8～ 太多了会乱哦～"
        )
    else:
        轮次 = 模板['epochs']
        图片尺寸 = 模板['imgsz']
        批量 = 模板['batch']
        早停耐心 = 模板['patience']
        线程数 = 模板['workers']
        print(f"\n{圈圈} 用「{模板['name']}」模式的参数哦～ 不用再设置啦～")
        print(f"轮次：{轮次} | 图片尺寸：{图片尺寸} | 批量：{批量} | 早停耐心：{早停耐心} | 线程数：{线程数}")

    # 设备选择（根据GPU可用性自动限制选项）
    print("\n设备选择说明（输?看帮助）：")
    if gpu可用:
        print("○ 单GPU训练：输编号（比如0）（快很多呢～）")
        print("○ 多GPU训练：输编号列表（比如0,1）（更快啦～）")
        print("○ 自动选闲置GPU：输 '-1'（单卡）或 '-1,-1'（多卡）")
        print("○ CPU训练：输 'cpu' （比较慢哦～）")
        if torch.backends.mps.is_available():
            print("○ 苹果电脑：输 'mps'")
    else:
        print("○ 没有可用GPU哦～ 只能选CPU训练呢～")
        print("○ CPU训练：输 'cpu'")

    设备 = None
    while True:
        设备输入 = input("要选哪个设备呢：").strip()
        if 设备输入 == "?":
            if gpu可用:
                print("\n？ 帮助：有GPU的话选GPU训练会快很多哦～ 输入编号就行啦～\n")
            else:
                print("\n？ 帮助：没有找到GPU呢～ 只能输入'cpu'用CPU训练哦～\n")
            continue
        
        # 无GPU时强制使用CPU
        if not gpu可用:
            if not 设备输入 or 设备输入.lower() != 'cpu':
                print(f"{叉叉} 没有可用的GPU哦～ 请输入'cpu'用CPU训练呢～")
                continue
            设备 = 'cpu'
            print(f"{星星} 选了CPU呢～ 虽然慢但很稳哦～")
            break
        
        # 有GPU时的处理
        if not 设备输入:
            设备 = None
            print(f"{星星} 自动选可用的GPU啦～ 真聪明～")
            break
        if ',' in 设备输入:
            try:
                设备列表 = [int(p.strip()) for p in 设备输入.split(',')]
                if all(d < torch.cuda.device_count() for d in 设备列表):
                    设备 = 设备列表
                    print(f"{星星} 选了这些GPU：{设备} 真棒～")
                    break
                else:
                    print(f"{叉叉} GPU编号不对哦～ 只有{torch.cuda.device_count()}个GPU呢～")
            except ValueError:
                print(f"{叉叉} 格式不对哦～ 应该像0,1这样的～ 再试试嘛～")
        else:
            if 设备输入.lower() == 'cpu':
                设备 = 'cpu'
                print(f"{星星} 选了CPU呢～ 虽然慢但很稳哦～")
                break
            if 设备输入.lower() == 'mps' and torch.backends.mps.is_available():
                设备 = 'mps'
                print(f"{星星} 选了苹果硅MPS～ 真厉害～")
                break
            try:
                设备编号 = int(设备输入)
                if 设备编号 < torch.cuda.device_count():
                    设备 = 设备编号
                    print(f"{星星} 选了GPU {设备编号}～ 速度肯定很快～")
                    break
                else:
                    print(f"{叉叉} GPU编号不对哦～ 只有{torch.cuda.device_count()}个GPU呢～")
            except ValueError:
                print(f"{叉叉} 输入不对哦～ 看看上面的说明再输呀～")

    # 数据增强参数
    print("\n===== 数据增强设置～ 好多选项可以选呢～ =====")
    用增强 = input("要自定义数据增强吗？(y/n，默认n)：").strip().lower() == 'y'
    增强参数 = {}
    if 用增强:
        增强参数 = 选择数据增强()
    else:
        print(f"{圈圈} 用默认的数据增强参数哦～ 不用自己调啦～")
 
 # AMP设置（新增选项：关闭AMP）
    print("\n===== AMP设置～ =====")
    print("○ AMP是自动混合精度训练，能加速训练但连不上Github会出问题～")
    amp选择 = input("要关闭AMP吗？(y/n，默认y)：").strip().lower()
    使用AMP = not (amp选择 == 'y' or amp选择 == '')  # 输入y则关闭（使用AMP=False）
    print(f"{星星} AMP设置：{'开启' if 使用AMP else '关闭'}")

    # 保存路径配置（父文件夹+时间子文件夹）
    保存帮助 = "训练结果会保存在这个父文件夹下，自动按时间创建子文件夹（比如20240807_153045）～"
    父文件夹 = 选保存父文件夹(保存帮助, 历史父文件夹)
    子文件夹 = 创建时间子文件夹(父文件夹)  # 按当前时间创建子文件夹

    # 日志配置
    print("\n===== 日志设置～ =====")
    用日志 = input("要开TensorBoard日志吗？(y/n，默认y)：").strip().lower() != 'n'

    # 参数确认
    print("\n===== 训练参数确认～ =====")
    print(f"模型路径：{模型路径}")
    print(f"数据集路径：{数据路径}")
    print(f"训练模式：{模板['name']} | 轮次：{轮次} | 图片尺寸：{图片尺寸}")
    print(f"批量大小：{批量} | 早停耐心值：{早停耐心} | 线程数：{线程数}")
    print(f"设备：{设备} | 父文件夹：{父文件夹} | 实际保存路径：{子文件夹}")
    print(f"数据增强：{'自定义' if 用增强 else '默认'} | TensorBoard：{'开' if 用日志 else '关'}")
    
    # 添加迁移学习策略信息
    策略名称 = {"1": "固定特征提取器", "2": "微调所有层", "3": "两阶段训练"}[迁移学习策略]
    print(f"迁移学习策略：{策略名称} (策略{迁移学习策略})")
    if 迁移学习策略 == "1":
        print(f"  - 冻结层数：{冻结层数}")
    elif 迁移学习策略 == "2":
        print(f"  - 微调学习率比例：{微调学习率比例}")
    else:
        print(f"  - 冻结层数：{冻结层数} | 解冻轮次：{解冻轮次} | 微调学习率比例：{微调学习率比例}")
    
    if 用增强 and 增强参数:
        print("数据增强参数：" + ", ".join([f"{k}={v}" for k, v in 增强参数.items()]))

    # 开始训练
    用户输入 = input("\n确认要开始训练吗？(y/n，默认y)：").strip().lower()
    if 用户输入 == 'n':
        print("训练取消啦～ 下次再来哦～ (｡･ω･｡)ﾉ♡")
        return

    # 保存配置（记录父文件夹路径，下次默认使用）
    要保存的配置 = {
        "dataset": str(数据路径),
        "last_model": str(模型路径) if 续训 else 历史配置.get("last_model", ""),
        "pretrained_model": str(模型路径) if not 续训 else 历史配置.get("pretrained_model", ""),
        "save_parent_dir": str(父文件夹)  # 保存父文件夹路径
    }
    保存配置(要保存的配置)

    # 训练过程
    try:
        print(f"\n{箭头} 开始训练咯～ 加油加油～ 结果会保存在：{子文件夹}")
        
        # 用户选择是否启用余弦学习率调度器
        启用余弦学习率 = input(f"{问号} 是否启用余弦学习率调度器？(y/n，默认y): ").strip().lower()
        use_cos_lr = 启用余弦学习率 != 'n' and 启用余弦学习率 != 'no'
        print(f"{圈圈} {'已启用' if use_cos_lr else '已禁用'} 余弦学习率调度器")
        
        # 设置训练参数
        训练参数 = {
            "data": str(数据路径),
            "epochs": 轮次,
            "imgsz": 图片尺寸,
            "batch": 批量,
            "patience": 早停耐心,
            "device": 设备,
            "project": str(父文件夹),  # 父文件夹
            "name": 子文件夹.name,  # 时间子文件夹名称
            "exist_ok": True,  # 允许创建同名时间子文件夹（实际不会重复，因为时间唯一）
            "workers": 线程数,
            "resume": 续训, **增强参数,
            "amp": 使用AMP,  # AMP参数控制
            "cos_lr": use_cos_lr,  # 根据用户选择启用或禁用余弦学习率调度器
            # 添加模型路径参数
            "model": str(模型路径),
        }
        
        # 添加迁移学习参数处理函数
        def 配置迁移学习(模型, 迁移学习策略, 冻结层数, 微调学习率比例, 解冻轮次, 训练参数):
            """自动配置迁移学习参数并应用到训练参数中"""
            
            # 添加freeze参数到训练参数中
            训练参数["freeze"] = 冻结层数 if 冻结层数 > 0 else None
            
            # 处理冻结层数为0或null的情况
            if 冻结层数 <= 0:
                # 确保所有层都可训练
                for param in 模型.model.parameters():
                    param.requires_grad = True
                print(f"{圈圈} 未冻结任何层，所有层都将参与训练～")
                
                if 迁移学习策略 == "1":
                    # 策略1但不冻结任何层，相当于普通训练
                    return 训练参数, None
                # 其他策略继续处理
            
            if 迁移学习策略 == "1" and 冻结层数 > 0:
                # 手动冻结前N层
                total_layers = len(list(模型.model.parameters()))
                for i, param in enumerate(模型.model.parameters()):
                    if i < 冻结层数:
                        param.requires_grad = False
                    else:
                        param.requires_grad = True
                print(f"{圈圈} 已手动冻结前{冻结层数}层特征提取层，只训练检测头层～")
                return 训练参数, None
                
            elif 迁移学习策略 == "2":
                print(f"{圈圈} 微调所有层，特征提取层学习率为检测头层的{微调学习率比例}倍～")
                
                # 改为设置优化器名称和学习率参数
                训练参数["optimizer"] = 'AdamW'
                训练参数["lr0"] = 训练参数.get('lr0', 0.001) * 微调学习率比例
                训练参数["lrf"] = 0.01  # 设置学习率衰减因子
                return 训练参数, None
            
            else:  # 策略3
                阶段参数 = []
            
            if 冻结层数 > 0:
                # 第一阶段参数
                阶段1参数 = 训练参数.copy()
                阶段1参数["epochs"] = 解冻轮次
                阶段1参数["cos_lr"] = use_cos_lr  # 使用用户选择的余弦学习率设置
                阶段1参数["resume"] = False  # 新增：明确设置冻结阶段resume为False
                # 添加lr0参数，移除可能存在的lr参数
                if 'lr' in 阶段1参数:
                    阶段1参数['lr0'] = 阶段1参数.pop('lr')
                阶段参数.append(("冻结阶段", 阶段1参数, 冻结层数))
                
                # 第二阶段参数
                阶段2参数 = 训练参数.copy()
                阶段2参数["epochs"] = 轮次 - 解冻轮次  # 正确设置微调阶段epochs
                # 设置model参数为第一阶段的last.pt
                阶段2模型路径 = str(Path(阶段1参数["project"]) / 阶段1参数["name"] / "weights" / "last.pt")
                # 打印路径以验证
                print(f"{箭头} 微调阶段模型路径: {阶段2模型路径}")
                阶段2参数["resume"] = False  # 修改：设置微调阶段resume为False
                阶段2参数["model"] = 阶段2模型路径  # 确保使用上一阶段的last.pt
                阶段2参数["cos_lr"] = use_cos_lr  # 使用用户选择的余弦学习率设置
                # 添加lr0参数，移除可能存在的lr参数
                if 'lr' in 阶段2参数:
                    阶段2参数['lr0'] = 阶段2参数.pop('lr')
                # 确保微调阶段使用较小的学习率
                阶段2参数['lr0'] = 阶段2参数.get('lr0', 0.001) * 微调学习率比例
                阶段参数.append(("微调阶段", 阶段2参数, 0, 微调学习率比例))
            else:
                # 冻结层数为0，直接进行单阶段训练
                print(f"{圈圈} 策略3但未冻结任何层，将进行单阶段训练～")
                # 直接返回训练参数和空的阶段参数
                return 训练参数, None
                
            return 训练参数, 阶段参数
        
        # 应用迁移学习配置
        训练参数, 阶段参数 = 配置迁移学习(模型, 迁移学习策略, 冻结层数, 微调学习率比例, 解冻轮次, 训练参数)
        
        # 打印训练参数以验证
        print(f"{箭头} 训练参数: {训练参数}")

        # 执行训练
        结果 = None
        if 阶段参数 and 迁移学习策略 == "3":
            # 策略3：两阶段训练特殊处理
            阶段名称, 阶段1参数, 冻结层数 = 阶段参数[0]
            print(f"{箭头} 开始{阶段名称}，使用模型: {阶段1参数.get('model', '未设置')}")
            print(f"{箭头} 阶段参数: {阶段1参数}")
            
            # 手动冻结前N层
            total_layers = len(list(模型.model.parameters()))
            for i, param in enumerate(模型.model.parameters()):
                if i < 冻结层数:
                    param.requires_grad = False
                else:
                    param.requires_grad = True
            print(f"{圈圈} {阶段名称}：已手动冻结前{冻结层数}层特征提取层～")
            
            # 执行冻结阶段训练
            冻结阶段结果 = 模型.train(**阶段1参数)
            
            # 保存冻结阶段的last.pt路径和学习率
            冻结阶段_last_pt = str(Path(阶段1参数["project"]) / 阶段1参数["name"] / "weights" / "last.pt")
            冻结阶段_lr = 阶段1参数.get('lr0', 0.001)
            print(f"{箭头} 冻结阶段训练完成，模型保存路径: {冻结阶段_last_pt}")
            print(f"{箭头} 冻结阶段学习率: {冻结阶段_lr}")
            
            # 询问是否开启解冻训练，等待10秒
            print(f"\n{问号} 是否开启解冻训练？(y/n，10秒后无输入将自动开始)：")
            import time, sys
            start_time = time.time()
            解冻训练_input = None
            
            # 检测操作系统
            import platform
            is_windows = platform.system() == 'Windows'
            
            if is_windows:
                # Windows系统使用msvcrt模块
                import msvcrt
                while time.time() - start_time < 10:
                    if msvcrt.kbhit():
                        解冻训练_input = msvcrt.getche().lower()
                        # 读取剩余的换行符
                        if msvcrt.kbhit() and msvcrt.getch() == b'\r':
                            pass
                        break
                    time.sleep(0.1)
                # 如果有输入且是'y'或'\n'，则视为同意
                if 解冻训练_input is not None:
                    解冻训练_input = 解冻训练_input.decode('utf-8') if isinstance(解冻训练_input, bytes) else 解冻训练_input
            else:
                # 非Windows系统使用select
                import select
                while time.time() - start_time < 10:
                    if select.select([sys.stdin], [], [], 0)[0]:
                        解冻训练_input = sys.stdin.readline().strip().lower()
                        break
                    time.sleep(0.1)
            
            # 判断是否开启解冻训练
            if 解冻训练_input is None or 解冻训练_input in ['', 'y', 'yes']:
                print(f"{圈圈} 开始解冻训练～")
                
                # 创建新的按时间命名文件夹_Fine-tuning
                当前时间 = datetime.now().strftime("%Y%m%d_%H%M%S")
                微调子文件夹_name = f"{当前时间}_Fine-tuning"
                
                # 设置解冻训练参数
                微调参数 = 训练参数.copy()
                微调参数["epochs"] = 轮次 - 解冻轮次  # 总轮数减去冻结轮次
                微调参数["model"] = 冻结阶段_last_pt  # 使用冻结阶段的last.pt
                微调参数["name"] = 微调子文件夹_name  # 新的时间命名文件夹
                微调参数["resume"] = False  # 解冻训练视为新训练
                微调参数["cos_lr"] = True  # 强制开启余弦学习率调度器
                微调参数["lr0"] = 冻结阶段_lr * 微调学习率比例  # 使用更小的学习率（修复：将lr改为lr0）
                微调参数["freeze"] = None  # 解冻所有层
                
                print(f"{箭头} 解冻训练参数: {微调参数}")
                print(f"{箭头} 解冻训练结果将保存在: {Path(微调参数['project']) / 微调参数['name']}")
                
                # 解冻所有层
                for param in 模型.model.parameters():
                    param.requires_grad = True
                print(f"{星星} 已解冻所有层，使用不同学习率进行微调～")
                
                # 获取模型参数组并设置不同学习率
                特征提取参数 = []
                检测头参数 = []
                total_layers = len(list(模型.model.parameters()))
                feature_layers = int(total_layers * 0.8)
                
                for i, (name, param) in enumerate(模型.model.named_parameters()):
                    if i < feature_layers:
                        特征提取参数.append(param)
                    else:
                        检测头参数.append(param)
                
                # 自定义优化器
                base_lr = 微调参数.get('lr0', 0.001)  # 修复：将lr改为lr0
                # 移除直接创建优化器对象的代码
                # 优化器 = torch.optim.AdamW([
                #     {'params': 特征提取参数, 'lr': base_lr * 微调学习率比例},
                #     {'params': 检测头参数, 'lr': base_lr}
                # ])
                # 微调参数["optimizer"] = 优化器
                
                # 改为设置优化器名称和学习率参数
                微调参数["optimizer"] = 'AdamW'
                微调参数["lr0"] = base_lr
                微调参数["lrf"] = 0.01  # 设置学习率衰减因子
                
                # 执行解冻训练
                解冻阶段结果 = 模型.train(**微调参数)
                print(f"{星星} 解冻阶段训练完成！")
                
                # 合并结果
                结果 = {"冻结阶段": 冻结阶段结果, "解冻阶段": 解冻阶段结果}
            else:
                print(f"{圈圈} 已取消解冻训练～")
                结果 = {"冻结阶段": 冻结阶段结果}
        elif 阶段参数:
            # 其他多阶段训练逻辑保持不变
            for 阶段名称, 阶段参数, *额外参数 in 阶段参数:
                print(f"{箭头} 开始{阶段名称}，使用模型: {阶段参数.get('model', '未设置')}")
                print(f"{箭头} 阶段参数: {阶段参数}")
                if 阶段名称 == "冻结阶段":
                    冻结层数 = 额外参数[0]
                    # 手动冻结前N层
                    total_layers = len(list(模型.model.parameters()))
                    for i, param in enumerate(模型.model.parameters()):
                        if i < 冻结层数:
                            param.requires_grad = False
                        else:
                            param.requires_grad = True
                    print(f"{圈圈} {阶段名称}：已手动冻结前{冻结层数}层特征提取层～")
                else:
                        # 手动解冻所有层
                        for param in 模型.model.parameters():
                            param.requires_grad = True
                        微调学习率比例 = 额外参数[0]
                        print(f"{星星} {阶段名称}：已解冻所有层，使用不同学习率进行微调～")
                        
                        # 获取模型参数组
                        特征提取参数 = []
                        检测头参数 = []
                        total_layers = len(list(模型.model.parameters()))
                        feature_layers = int(total_layers * 0.8)
                          
                        for i, (name, param) in enumerate(模型.model.named_parameters()):
                            if i < feature_layers:
                                特征提取参数.append(param)
                            else:
                                检测头参数.append(param)
                          
                        # 自定义优化器
                        base_lr = 阶段参数.get('lr0', 0.001)  # 修复：将lr改为lr0
                        # 移除直接创建优化器对象的代码
                        # 优化器 = torch.optim.AdamW([
                        #     {'params': 特征提取参数, 'lr': base_lr * 微调学习率比例},
                        #     {'params': 检测头参数, 'lr': base_lr}
                        # ])
                        # 阶段参数["optimizer"] = 优化器
                        
                        # 改为设置优化器名称和学习率参数
                        阶段参数["optimizer"] = 'AdamW'
                        阶段参数["lr0"] = base_lr
                        阶段参数["lrf"] = 0.01  # 设置学习率衰减因子
                
                结果 = 模型.train(**阶段参数)
        else:
            # 单阶段训练
            结果 = 模型.train(**训练参数)
        
        # 模型评估
        print(f"\n{箭头} 开始评估模型性能～")
        评估结果 = 模型.val(data=str(数据路径), device=设备)
        
        print(f"\n{星星} 训练和评估完成啦～ 太棒啦～ 结果保存在：{子文件夹}")
        
        # 修复：根据结果类型正确显示训练指标
        if isinstance(结果, dict):
            # 多阶段训练结果
            for 阶段名称, 阶段结果 in 结果.items():
                if hasattr(阶段结果, 'box'):
                    print(f"☆ {阶段名称}训练指标：mAP50={阶段结果.box.map50:.3f} | mAP50-95={阶段结果.box.map:.3f}")
                else:
                    print(f"☆ {阶段名称}训练完成，但未找到性能指标")
        else:
            # 单阶段训练结果
            if hasattr(结果, 'box'):
                print(f"☆ 训练指标：mAP50={结果.box.map50:.3f} | mAP50-95={结果.box.map:.3f}")
            else:
                print("☆ 训练完成，但未找到性能指标")
        
        # 评估结果应该始终是YOLO结果对象
        if hasattr(评估结果, 'box'):
            print(f"☆ 评估指标：mAP50={评估结果.box.map50:.3f} | mAP50-95={评估结果.box.map:.3f}")
        else:
            print("☆ 评估完成，但未找到性能指标")
    except Exception as e:
        print(f"{叉叉} 训练出了点问题：{str(e)} 别难过，再试试呀～")

if __name__ == "__main__":
    main()